from __future__ import annotations

import mmap
import os
import struct
import sys


''' DevMem
Class to read and write data aligned to word boundaries of /dev/mem
'''


class DevMem:
    __slots__ = ['address', 'start_offset', 'word_size', 'number_of_words', 'word_format', 'file', 'mem', 'data']

    word_formats = {1: 'B', 2: 'H', 4: 'I'}

    def __init__(self, address: int, number_of_words: int = 1, word_size: int = 4, filename: str = '/dev/mem') -> None:
        if address < 0 or number_of_words < 1 or word_size not in self.word_formats.keys():
            raise AssertionError
        if not isinstance(address, int) or not isinstance(word_size, int) or not isinstance(filename, str):
            raise TypeError
        self.address = address & -mmap.PAGESIZE
        self.start_offset = address - self.address
        self.word_size = word_size
        self.number_of_words = number_of_words
        self.word_format = self.word_formats.get(word_size)
        self.file = os.open(filename, os.O_RDWR | os.O_SYNC)
        end_offset = address + (number_of_words + 1) * word_size
        if end_offset % self.word_size:
            raise AssertionError
        self.mem = mmap.mmap(self.file, end_offset - self.address, mmap.MAP_SHARED,
                             mmap.PROT_READ | mmap.PROT_WRITE, offset=self.address)
        self.mem.seek(self.start_offset)
        self.data = self.mem.read(self.word_size)

    def read(self, offset: int = 0) -> DevMem:
        if offset < 0 or offset % self.word_size:
            raise AssertionError
        self.mem.seek(self.start_offset + offset)
        self.data = self.mem.read(self.word_size)
        return self

    def write(self, offset: int = 0, value: int = 0) -> DevMem:
        if offset < 0 or offset % self.word_size:
            raise AssertionError
        value = round(value)
        if value < 0:
            value = (value + (1 << self.word_size * 8)) % (1 << self.word_size * 8)
        self.mem.seek(self.start_offset + offset)
        self.mem.write(struct.pack(self.word_formats.get(self.word_size), value))
        self.mem.seek(self.start_offset + offset)
        self.data = self.mem.read(self.word_size)
        return self

    def get_value(self, signed=False) -> int:
        value = int.from_bytes(self.data, byteorder=sys.byteorder, signed=signed)
        return value

    def __getitem__(self, i: int) -> bool:
        if not isinstance(i, int):
            raise AssertionError
        return True if int.from_bytes(self.data, byteorder=sys.byteorder) & 1 << i != 0 else False

    def __setitem__(self, i: int, value: bool) -> None:
        if not isinstance(i, int) and not isinstance(value, bool):
            raise AssertionError
        mask = 1 << i
        reg_value = int.from_bytes(self.data, byteorder=sys.byteorder)
        reg_value &= ~mask
        if value:
            reg_value |= mask
        self.write(offset=0, value=reg_value)

    def __str__(self) -> str:
        return hex(int.from_bytes(self.data, byteorder=sys.byteorder))

    def __del__(self) -> None:
        if self.file:
            os.close(self.file)


class DevMemCache:
    __slots__ = ['_map', '_number_of_words']

    def __init__(self, number_of_words: int = 0x100):
        self._map: dict[(int, int), DevMem] = dict()
        self._number_of_words = number_of_words

    def get(self, address: int, word_size: int = 4) -> DevMem:
        key = (address, word_size)
        return self._map.get(key, DevMem(address=address, word_size=word_size, number_of_words=self._number_of_words))
